"""
Basic usage example of the Claude Code Theory framework.

This example demonstrates:
1. Setting up a Claude Code system
2. Creating tools and action spaces  
3. Implementing contextual bandits for tool selection
4. Running a simple optimization loop
"""

using Pkg
Pkg.activate(".")

# Include the package
include("../src/ClaudeCodeTheory.jl")
using .ClaudeCodeTheory

using LinearAlgebra
using Random
using Statistics
using Printf

# ============================================================================
# Example 1: Basic Tool Definition and Action Space
# ============================================================================

function example_tool_setup()
    println("=== Example 1: Tool Setup ===")
    
    # Define some example tools
    tools = [
        Tool(
            "read_file", FILE_TOOLS, "Read contents of a file",
            Dict("path" => String), Dict("content" => String),
            SAFE, 1.0, 0.5, true, Set{String}()
        ),
        Tool(
            "write_file", FILE_TOOLS, "Write content to a file", 
            Dict("path" => String, "content" => String), Dict("success" => Bool),
            MEDIUM_RISK, 2.0, 1.0, false, Set{String}()
        ),
        Tool(
            "search_code", SEARCH_TOOLS, "Search for code patterns",
            Dict("pattern" => String, "directory" => String), Dict("matches" => Vector),
            SAFE, 1.5, 2.0, true, Set{String}()
        ),
        Tool(
            "run_tests", TEST_TOOLS, "Execute test suite",
            Dict("test_path" => String), Dict("results" => Dict),
            LOW_RISK, 5.0, 10.0, false, Set{String}()
        ),
        Tool(
            "git_commit", GIT_TOOLS, "Create git commit",
            Dict("message" => String, "files" => Vector), Dict("commit_hash" => String),
            MEDIUM_RISK, 3.0, 3.0, false, Set{String}()
        )
    ]
    
    # Create action space
    action_space = ClaudeCodeActionSpace(tools, 3)  # max 3 concurrent tools
    
    println("Created action space with $(length(tools)) tools")
    println("Tool compatibility matrix:")
    display(action_space.compatibility_matrix)
    
    # Test action creation
    params = Dict{String, Any}("path" => "/tmp/test.txt")
    action = create_action(action_space, "read_file", params, 1.0, 30.0)
    
    if action !== nothing
        println("\nCreated action: $(action.tool.name)")
        println("Parameters: $(action.parameters)")
        println("Estimated cost: $(estimate_action_cost(action))")
        println("Estimated latency: $(estimate_action_latency(action))")
        println("Risk level: $(action_risk_level(action))")
    end
    
    return action_space
end

# ============================================================================
# Example 2: Contextual Bandit Tool Selection
# ============================================================================

function example_contextual_bandits()
    println("\n=== Example 2: Contextual Bandit Tool Selection ===")
    
    # Setup
    d = 10  # feature dimension
    k = 5   # number of tools (arms)
    
    # Create bandit agents
    linucb_agent = LinUCBAgent(d, k, 1.0)
    thompson_agent = ThompsonSamplingAgent(d, k)
    
    # Create feature extractor
    feature_extractor = ContextFeatureExtractor(d)
    
    println("Created LinUCB agent with regret bound: $(get_regret_bound(linucb_agent, 1000))")
    println("Created Thompson Sampling agent with regret bound: $(get_regret_bound(thompson_agent, 1000))")
    
    # Simulate some tool selection rounds
    Random.seed!(42)
    num_rounds = 100
    
    println("\nRunning $num_rounds rounds of tool selection...")
    
    linucb_rewards = Float64[]
    thompson_rewards = Float64[]
    
    # Simulate true reward function (unknown to agents)
    true_theta = [randn(d) for _ in 1:k]  # true parameters for each tool
    
    for round in 1:num_rounds
        # Generate random context
        context = randn(d)
        context = context / norm(context)  # normalize
        
        # LinUCB selection
        linucb_arm, linucb_confidence = select_arm(linucb_agent, context)
        linucb_reward = dot(context, true_theta[linucb_arm]) + 0.1 * randn()  # noisy reward
        update_bandit!(linucb_agent, linucb_arm, linucb_reward, context)
        push!(linucb_rewards, linucb_reward)
        
        # Thompson Sampling selection  
        thompson_arm, thompson_uncertainty = select_arm(thompson_agent, context)
        thompson_reward = dot(context, true_theta[thompson_arm]) + 0.1 * randn()  # noisy reward
        update_bandit!(thompson_agent, thompson_arm, thompson_reward, context)
        push!(thompson_rewards, thompson_reward)
        
        # Update feature normalization
        update_normalization!(feature_extractor, context)
        
        if round % 20 == 0
            println("Round $round: LinUCB avg reward = $(mean(linucb_rewards[max(1, round-19):round]):.3f), " *
                   "Thompson avg reward = $(mean(thompson_rewards[max(1, round-19):round]):.3f)")
        end
    end
    
    # Final performance analysis
    println("\nFinal Performance:")
    println("LinUCB average reward: $(mean(linucb_rewards):.3f)")
    println("Thompson Sampling average reward: $(mean(thompson_rewards):.3f)")
    
    # Analyze selection patterns
    linucb_analysis = analyze_arm_selection_pattern(linucb_agent)
    thompson_analysis = analyze_arm_selection_pattern(thompson_agent)
    
    println("\nExploration Analysis:")
    println("LinUCB exploration rate: $(linucb_analysis.exploration_rate:.3f)")
    println("Thompson exploration rate: $(thompson_analysis.exploration_rate:.3f)")
    
    return linucb_agent, thompson_agent
end

# ============================================================================
# Example 3: Multi-Objective Utility Optimization  
# ============================================================================

function example_multi_objective_optimization()
    println("\n=== Example 3: Multi-Objective Optimization ===")
    
    # Define multiple objectives with different weights
    scenarios = [
        ("Quality-focused", [0.6, 0.2, 0.1, 0.1]),   # emphasize quality
        ("Speed-focused", [0.3, 0.5, 0.1, 0.1]),     # emphasize low latency  
        ("Cost-conscious", [0.3, 0.2, 0.4, 0.1]),    # emphasize low cost
        ("Risk-averse", [0.3, 0.2, 0.2, 0.3])        # emphasize low risk
    ]
    
    # Simulate different actions with varying quality/latency/cost/risk profiles
    actions_data = [
        ("quick_fix", 0.6, 1.0, 2.0, 0.3),      # low quality, fast, cheap, medium risk
        ("thorough_solution", 0.9, 5.0, 8.0, 0.1),  # high quality, slow, expensive, low risk
        ("automated_tool", 0.7, 2.0, 1.0, 0.2),     # good quality, medium speed, cheap, low risk
        ("manual_process", 0.8, 10.0, 15.0, 0.05),  # high quality, very slow, expensive, very low risk
        ("experimental_approach", 0.5, 3.0, 5.0, 0.6)  # uncertain quality, medium speed, medium cost, high risk
    ]
    
    println("Comparing utility across different scenarios:")
    println("Action".rjust(20) * " | " * join([scenario[1].rjust(15) for scenario in scenarios], " | "))
    println("-" ^ (20 + 4 * 16 + 3))
    
    for (action_name, quality, latency, cost, risk) in actions_data
        utilities = []
        for (scenario_name, weights) in scenarios
            utility = MultiObjectiveUtility(weights, quality, latency, cost, risk)
            push!(utilities, utility_value(utility))
        end
        
        utility_str = join([@sprintf("%.3f", u) for u in utilities], " | ")
        println(action_name.rjust(20) * " | " * utility_str)
    end
    
    # Find Pareto frontier
    objectives = [[data[2], -data[3], -data[4], -data[5]] for data in actions_data]  # maximize quality, minimize others
    pareto_indices = pareto_frontier(objectives)
    
    println("\nPareto optimal actions:")
    for idx in pareto_indices
        println("  - $(actions_data[idx][1])")
    end
end

# Helper function for string formatting (no longer needed)

# ============================================================================
# Example 4: State Space and Environment Modeling
# ============================================================================

function example_state_space()
    println("\n=== Example 4: State Space Modeling ===")
    
    # Create initial state from current directory (if it exists and is safe)
    current_dir = pwd()
    println("Analyzing current directory: $current_dir")
    
    try
        # Create initial environmental state
        state = create_initial_state(current_dir)
        
        println("State created successfully:")
        println("  Files: $(length(state.filesystem.files))")
        println("  Directories: $(length(state.filesystem.directories))")
        println("  Dependency nodes: $(length(state.dependencies.nodes))")
        println("  Dependency edges: $(length(state.dependencies.edges))")
        println("  Available tools: $(length(state.tools.available_tools))")
        println("  System load: $(state.external.system_load)")
        println("  Network connectivity: $(state.external.network_connectivity)")
        
        # Create a second state to test distance computation
        state2 = create_initial_state(current_dir)
        distance = state_distance(state, state2)
        println("  Distance between identical states: $distance")
        
        return state
        
    catch e
        println("Could not analyze directory: $e")
        return nothing
    end
end

# ============================================================================
# Main Example Runner
# ============================================================================

function run_examples()
    println("Claude Code Theory Framework - Basic Examples")
    println("=" ^ 60)
    
    try
        # Run examples
        action_space = example_tool_setup()
        linucb_agent, thompson_agent = example_contextual_bandits()
        example_multi_objective_optimization()
        state = example_state_space()
        
        println("\n" * "=" ^ 60)
        println("All examples completed successfully!")
        
        # Optional: demonstrate serialization
        println("\n=== Bonus: Agent Serialization ===")
        agent_state = export_agent_state(linucb_agent)
        println("Exported LinUCB agent state with $(length(agent_state)) fields")
        
        # Create new agent and import state
        new_agent = LinUCBAgent(10, 5)
        import_agent_state!(new_agent, agent_state)
        println("Successfully imported agent state")
        
    catch e
        println("\nError running examples: $e")
        println("Stack trace:")
        for frame in stacktrace(catch_backtrace())
            println("  $frame")
        end
    end
end

# Run the examples if this script is executed directly
if abspath(PROGRAM_FILE) == @__FILE__
    run_examples()
end